// SPDX-FileCopyrightText: 2021 - 2022 UnionTech Software Technology Co., Ltd.
//
// SPDX-License-Identifier: GPL-3.0-or-later

#include "rsahelper.h"

#include <QDebug>

#define BEGIN_RSA_PUBLIC_KEY  "BEGIN RSA PUBLIC KEY"
#define BEGIN_RSA_PRIVATE_KEY  "BEGIN RSA PRIVATE KEY"
#define BEGIN_PUBLIC_KEY      "BEGIN PUBLIC KEY"
#define BEGIN_PRIVATE_KEY      "BEGIN PRIVATE KEY"

const int L1024PublicKeyTypeUnitStrLength = 117;
const int L2048PublicKeyTypeUnitStrLength = 245;
const int L1024PrivateKeyTypeUnitStrLength = 172;
const int L2048PrivateKeyTypeUnitStrLength = 344;

QStringList spiltString(const QString& data, int n) {
    QStringList res;
    QString tempStr = data;
    while (!tempStr.isEmpty()) {
        res.append(tempStr.left(n));
        tempStr = tempStr.mid(n);
    }
    return res;
}

RsaHelper::RsaHelper(QObject *parent)
    : QObject(parent)
    , m_keyType(L2048)
    , m_publicKeyBio(nullptr)
    , m_publicKeyRsa(nullptr)
    , m_privateKeyBio(nullptr)
    , m_privateKeyRsa(nullptr)
{
}

RsaHelper::~RsaHelper()
{
    if (m_publicKeyBio) {
        BIO_free_all(m_publicKeyBio);
    }
    if (m_publicKeyRsa) {
        RSA_free(m_publicKeyRsa);
    }

    if (m_privateKeyBio) {
        BIO_free_all(m_privateKeyBio);
    }
    if (m_privateKeyRsa) {
        RSA_free(m_privateKeyRsa);
    }
}

void RsaHelper::setPublicKey(const QString &key)
{
    if (m_publicKeyBio) {
        BIO_free_all(m_publicKeyBio);
        m_publicKeyBio = nullptr;
    }
    if (m_publicKeyRsa) {
        RSA_free(m_publicKeyRsa);
        m_publicKeyRsa = nullptr;
    }

    const QByteArray &keyByteArray = key.toUtf8();

    m_publicKeyBio = BIO_new_mem_buf(keyByteArray.data(), keyByteArray.length());
    if (nullptr == m_publicKeyBio) {
        qWarning() << Q_FUNC_INFO << "new a bio failed";
        return;
    }

    // 从BIO中读取公钥对象
    m_publicKeyRsa = RSA_new();
    if (key.contains(BEGIN_RSA_PUBLIC_KEY)) {
        m_publicKeyRsa = PEM_read_bio_RSAPublicKey(m_publicKeyBio, &m_publicKeyRsa, nullptr, nullptr);
    } else {
        m_publicKeyRsa = PEM_read_bio_RSA_PUBKEY(m_publicKeyBio, &m_publicKeyRsa, nullptr, nullptr);
    }
    if (nullptr == m_publicKeyRsa) {
        qWarning() << Q_FUNC_INFO << "rsa created failed";
        BIO_free_all(m_publicKeyBio);
        m_publicKeyBio = nullptr;
        return;
    }
}

void RsaHelper::setPrivateKey(const QString &key)
{
    if (m_privateKeyBio) {
        BIO_free_all(m_privateKeyBio);
        m_privateKeyBio = nullptr;
    }
    if (m_privateKeyRsa) {
        RSA_free(m_privateKeyRsa);
        m_privateKeyRsa = nullptr;
    }

    const QByteArray &keyByteArray = key.toUtf8();
    m_privateKeyBio = BIO_new_mem_buf(keyByteArray.data(), keyByteArray.length());
    if (nullptr == m_privateKeyBio) {
        qWarning() << Q_FUNC_INFO << "input private key is error";
        return;
    }

    m_privateKeyRsa = RSA_new();
    m_privateKeyRsa = PEM_read_bio_RSAPrivateKey(m_privateKeyBio, &m_privateKeyRsa, nullptr, nullptr);
    if (nullptr == m_privateKeyRsa) {
        qWarning() << Q_FUNC_INFO << "rsa created failed";
        BIO_free_all(m_privateKeyBio);
        m_privateKeyBio = nullptr;
        return;
    }
}

QString RsaHelper::encrypt(const QString &data)
{
    if (nullptr == m_publicKeyRsa) {
        qCritical() << Q_FUNC_INFO << "no valid public key rsa";
        return data;
    }

    int unitStrLength = 0;
    switch (m_keyType) {
    case L1024:
        unitStrLength = L1024PublicKeyTypeUnitStrLength;
        break;
    case L2048:
        unitStrLength = L2048PublicKeyTypeUnitStrLength;
        break;
    }

    QStringList dataUnitStrList = spiltString(data, unitStrLength);
    for (QString &unitStr : dataUnitStrList) {
        unsigned long nLen = static_cast<unsigned long>(RSA_size(m_publicKeyRsa));
        char *encryptBuf = new char[nLen];
        memset(encryptBuf, 0, nLen);
        QByteArray unitByteArry = unitStr.toUtf8();
        int dataCharPtrLen = unitByteArry.length();
        const uchar *dataCharPtr = reinterpret_cast<uchar *>(unitByteArry.data());
        int nSize = RSA_public_encrypt(dataCharPtrLen,
                                       dataCharPtr,
                                       reinterpret_cast<unsigned char *>(encryptBuf),
                                       m_publicKeyRsa,
                                       RSA_PKCS1_PADDING);
        if (0 > nSize) {
            return "";
        }
        QByteArray ba(encryptBuf, nSize);

        unitStr = ba.toBase64();
        delete[] encryptBuf;
    }

    return dataUnitStrList.join("");
}

QString RsaHelper::decrypt(const QString &data)
{
    if (nullptr == m_privateKeyRsa) {
        qCritical() << Q_FUNC_INFO << "no valid private key rsa";
        return data;
    }

    int unitStrLength = 0;
    switch (m_keyType) {
    case L1024:
        unitStrLength = L1024PrivateKeyTypeUnitStrLength;
        break;
    case L2048:
        unitStrLength = L2048PrivateKeyTypeUnitStrLength;
        break;
    }

    QStringList dataUnitStrList = spiltString(data, unitStrLength);
    for (QString &unitStr : dataUnitStrList) {
        unsigned long nLen = static_cast<unsigned long>(RSA_size(m_privateKeyRsa));
        char *decryptBuf = new char[nLen];
        memset(decryptBuf, 0, nLen);
        // 解base64
        QByteArray baFromBase64 = QByteArray::fromBase64(unitStr.toUtf8());
        const uchar *dataCharPtr = reinterpret_cast<uchar *>(baFromBase64.data());
        int dataCharPtrLen = baFromBase64.length();
        int nSize = RSA_private_decrypt(dataCharPtrLen,
                                       dataCharPtr,
                                       reinterpret_cast<uchar *>(decryptBuf),
                                       m_privateKeyRsa,
                                       RSA_PKCS1_PADDING);
        if (0 > nSize) {
            return "";
        }
        QByteArray ba(decryptBuf, nSize);

        unitStr = ba;
        delete[] decryptBuf;
    }
    return dataUnitStrList.join("");

}
